#include "stdafx.h"
#include "Model.h"

WMReaderModel::WMReaderModel()
{
    m_RapiInited = false;
    CoInitialize(NULL);
}

WMReaderModel::~WMReaderModel()
{
    CoUninitialize();

    if(m_RapiInited == true)
        CeRapiUninit();
} 

void WMReaderModel::CheckConnection()
{
    HRESULT result = E_FAIL;
    bool initialized = false;

    RAPIINIT rapiInit;
    memset(&rapiInit, 0, sizeof(rapiInit));

    rapiInit.cbSize = sizeof(rapiInit);
    
    if(m_RapiInited == true) {
        m_RapiInited = false;
        CeRapiUninit();
    }

    result = CeRapiInitEx(&rapiInit);
    if(result == S_OK) {
        DWORD init = 0;
        initialized = true;

        init = WaitForSingleObject(rapiInit.heRapiInit, connect_set::connectionTimeout);
        if(init == WAIT_OBJECT_0) {
            result = rapiInit.hrRapiInit;
        } else if(init == WAIT_TIMEOUT) {
            result = HRESULT_FROM_WIN32(ERROR_TIMEOUT);
        }
    }
    
    if(initialized == true)
        CeRapiUninit();

    if(result != S_OK)
        throw DeviceIsNotConnected();
}

void WMReaderModel::ConnectToDevice()
{
    try {
        HRESULT result = S_OK;
        
        CheckConnection();    

        result = CeRapiInit();
        if(result != S_OK)
            throw ConnectFailed();
        
        m_RapiInited = true;
        
        if(!IsLibInstalled())
            InstallLibToDevice();              

        emit DeviceConnected();
    }
    catch(ConnectFailed &e) {
        emit RequestShowErrorMessage(e.what());
    }
    catch(DeviceIsNotConnected &e) {
        emit RequestShowErrorMessage(e.what());
        emit DeviceDisconnected();
    }
}

void WMReaderModel::DisconnectFromDevice()
{
    try {
        HRESULT result = S_OK;
        
        if (m_RapiInited == true) {

            result = CeRapiUninit();
            if(result != S_OK)
                throw DisconnectFailed();

            m_RapiInited = false;
        }
    
        emit DeviceDisconnected();
    }
    catch(DisconnectFailed &e) {
        emit RequestShowErrorMessage(e.what());
    }
}

void WMReaderModel::GetAppMgrPath(std::wstring &appMgrPath)
{
    try {
        HRESULT result = S_OK;

        std::auto_ptr <RegKeyHolder> appMng(new RegKeyHolder());

        result = appMng->OpenKey(HKEY_LOCAL_MACHINE, cab_set::appMgrKey);
        if(result != S_OK)
            throw RetrieveAppMgrPathFailed();
    
        BYTE buffer[MAX_PATH];
        memset(buffer, 0, MAX_PATH);

        result = appMng->GetKeyValue(cab_set::appMgrValue, buffer, MAX_PATH);
        if(result != S_OK)
            throw RetrieveAppMgrPathFailed();

        appMgrPath = reinterpret_cast<const wchar_t*>(buffer);
    }
    catch(RetrieveAppMgrPathFailed &e) {
        DisconnectFromDevice();
        emit RequestShowErrorMessage(e.what());
    }
}

void WMReaderModel::InstallLibToDevice()
{
    try {
        std::wstring appMgrPath;
        GetAppMgrPath(appMgrPath);

        appMgrPath = appMgrPath + cab_set::appMgrBin;

        DWORD result = 0;

        wchar_t buffer[MAX_PATH];
        memset(buffer, 0, MAX_PATH);

        result = GetCurrentDirectory(MAX_PATH, buffer);
        if(result == 0)
            throw LibInstallFailed();

        std::wstring cabIniPath = buffer;
        cabIniPath = cabIniPath + cab_set::cabIni;
        
        result = reinterpret_cast<DWORD>(ShellExecute(NULL,
                                                      NULL,
                                                      appMgrPath.c_str(),
                                                      cabIniPath.c_str(),
                                                      NULL,
                                                      SW_SHOWNORMAL));
        if(result < 32)
            throw LibInstallFailed();
    }
    catch(LibInstallFailed &e) {
        DisconnectFromDevice();
        emit RequestShowErrorMessage(e.what());
    }
}

bool WMReaderModel::IsLibInstalled()
{
    HRESULT result = S_OK;
    bool dllInstalled = false;

    std::auto_ptr <CeFindHolder> dll(new CeFindHolder());

    CE_FIND_DATA fileSystemItem;
    memset(&fileSystemItem, 0, sizeof(fileSystemItem));    

    result = dll->FirstFile(pathToDll, fileSystemItem);
    if(result == S_OK)
        dllInstalled = true;

    return dllInstalled;  
}

template <typename reciveType>
void WMReaderModel::RetrieveDataFromDevice(const std::wstring &input,
                                           const std::wstring &funcName,
                                           std::vector <reciveType> &reciveItems)
{
    HRESULT result = S_OK;
    BYTE *output = NULL;
    DWORD outputSize = 0;

    std::auto_ptr <RapiStreamHolder> rapiStream(new RapiStreamHolder());

    result = rapiStream->Invoke(pathToDll,
                                funcName,
                                input,
                                output,
                                outputSize);
    if(result != S_OK)
        throw RetrieveDataFailed();

    result = rapiStream->SetReadTimeout(connect_set::readTimeout);
    if(result != S_OK)
        throw RetrieveDataFailed();

    reciveType reciveItem;
    memset(&reciveItem, 0, sizeof(reciveItem));

    DWORD len = 0;
    DWORD bytes = 0;

    bool receivingData = true;

    do { 
        result = rapiStream->Stream()->Read(&len, sizeof(len), &bytes);
        if(result != S_OK )
            throw RetrieveDataFailed();

        if (len == connect_set::lastBuffer)
            break;

        result = rapiStream->Stream()->Read(&reciveItem, len, &bytes);
        if(result != S_OK)
            throw RetrieveDataFailed();

        reciveItems.push_back(reciveItem);

        memset(&reciveItem, 0, sizeof(reciveItem));
        bytes = 0;
    } while(receivingData);  
}

void WMReaderModel::RetrieveFileSystem(QTreeWidgetItem *topItem)
{   
    try {
        QString path;
        WMReaderUtils::PathFromTreeItem(topItem, path);

        std::vector <CE_FIND_DATA> fileSystemItems;
        RetrieveDataFromDevice(path.toStdWString(), dll_func::fileSystem, fileSystemItems);

        emit FileSystemRetrieved(topItem, &fileSystemItems);
    }
    catch(RetrieveDataFailed &e) {
        try {
            CheckConnection();
            emit RequestShowErrorMessage(e.what());
        }
        catch(DeviceIsNotConnected &e) {
            emit RequestShowErrorMessage(e.what());
            emit DeviceDisconnected();
        }
    }
}


void WMReaderModel::RetrieveRegistryKeysList(QTreeWidgetItem *topItem)
{
    try {
        QString regKey;
        WMReaderUtils::RegKeyFromTreeItem(topItem, regKey);

        std::vector <registry_set::REGKEY> registryKeys;
        RetrieveDataFromDevice(regKey.toStdWString(), dll_func::registryKeys, registryKeys);

        emit RegistryKeysListRetrieved(topItem, &registryKeys);
    }
    catch(RetrieveDataFailed &e) {
        try {
            CheckConnection();
            emit RequestShowErrorMessage(e.what());
        }
        catch(DeviceIsNotConnected &e) {
            emit RequestShowErrorMessage(e.what());
            emit DeviceDisconnected();
        }   
    }
}

void WMReaderModel::RetrieveRegistryKeyValues(QTreeWidgetItem *topItem)
{
    try {
        QString regKey;
        WMReaderUtils::RegKeyFromTreeItem(topItem, regKey);

        std::vector <registry_set::REGKEYVALUE> keyValues;
        RetrieveDataFromDevice(regKey.toStdWString(), dll_func::registryValues, keyValues);

        emit RegistryKeyValuesRetrieved(&keyValues);
    }
    catch(RetrieveDataFailed &e) {
        try {
            CheckConnection();
            emit RequestShowErrorMessage(e.what());
        }
        catch(DeviceIsNotConnected &e) {
            emit RequestShowErrorMessage(e.what());
            emit DeviceDisconnected();
        }
    }
}

void WMReaderModel::RetrieveDeviceInfo()
{
    try {
        std::vector <device_info::DEVICEINFO> deviceInfo;
        RetrieveDataFromDevice(L"", dll_func::deviceInfo, deviceInfo);

        emit DeviceInfoRetrieved(&deviceInfo);
    }
    catch(RetrieveDataFailed &e) {
        try {
            CheckConnection();
            emit RequestShowErrorMessage(e.what());
        }
        catch(DeviceIsNotConnected &e) {
            emit RequestShowErrorMessage(e.what());
            emit DeviceDisconnected();
        }
    }
}